﻿/*
BF4Rcon.NET, provides a .NET implementation of BF4's RCON interface.
Copyright (C) 2013 Tim Marro

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace TimScripts.BF4Rcon.Frostbite
{
    public class PacketConverter : IPacketConverter<Packet>
    {
        private Queue<byte> Data = new Queue<byte>(16384);
        private Func<IEnumerable<byte>> RequestData;
        private Func<Task<IEnumerable<byte>>> RequestDataAsync;
        private const int NumWordsSize = sizeof(int);
        private const int WordSizeSize = sizeof(int);
        private const int HeaderSize = sizeof(int) + sizeof(int);
        private const int IsFromClientMask = (1 << 31);
        private const int IsResponseMask = (1 << 30);
        private const int SequenceMask = IsResponseMask - 1;

        public PacketConverter(Func<IEnumerable<byte>> dataRequestFunction, Func<Task<IEnumerable<byte>>> dataRequestFunctionAsync)
        {
            RequestData = dataRequestFunction;
            RequestDataAsync = dataRequestFunctionAsync;
        }

        public Packet GetPacket()
        {
            PacketHeader header;
            int numWords;
            byte[] wordBytes;
            IEnumerable<string> words;
            Packet result;

            EnsureBufferHasBytes(HeaderSize);

            header = DequeueHeader();

            EnsureBufferHasBytes(header.PacketSize - HeaderSize);

            numWords = DequeueInt();
            wordBytes = DequeueBlock(header.PacketSize - HeaderSize - NumWordsSize);
            words = new WordCollection(numWords, wordBytes);

            result = new Packet(header.IsFromClient, header.IsResponse, header.Sequence, words);

            return result;
        }

        public async Task<Packet> GetPacketAsync()
        {
            PacketHeader header;
            int numWords;
            byte[] wordBytes;
            IEnumerable<string> words;
            Packet result;

            await EnsureBufferHasBytesAsync(HeaderSize).ConfigureAwait(false);

            header = DequeueHeader();

            await EnsureBufferHasBytesAsync(header.PacketSize - HeaderSize).ConfigureAwait(false);

            numWords = DequeueInt();
            wordBytes = DequeueBlock(header.PacketSize - HeaderSize - NumWordsSize);
            words = new WordCollection(numWords, wordBytes);

            result = new Packet(header.IsFromClient, header.IsResponse, header.Sequence, words);

            return result;
        }

        public byte[] GetPacketBytes(Packet packet)
        {
            if (packet == null)
                throw new ArgumentNullException("packet");

            int packetLength = CalculatePacketLength(packet);
            byte[] result = new byte[packetLength];
            int header = 0;
            int offset = 0;

            if (packet.IsFromClient) header |= 1 << 31;
            if (packet.IsResponse) header |= 1 << 30;
            header |= packet.Sequence;

            AddBytes(result, ref offset, IntToBytes(header));

            AddBytes(result, ref offset, IntToBytes(packetLength));

            AddBytes(result, ref offset, IntToBytes(packet.Words.Count));

            foreach (string word in packet.Words)
            {
                AddBytes(result, ref offset, IntToBytes(word.Length));

                AddBytes(result, ref offset, Encoding.UTF8.GetBytes(word));
                result[offset++] = 0;
            }

            return result;
        }

        #region GetPacketBytes Helpers
        private static void AddBytes(byte[] result, ref int offset, byte[] data)
        {
            Copy(result, data, ref offset);
        }

        private static void Copy(byte[] result, byte[] part, ref int offset)
        {
            Buffer.BlockCopy(part, 0, result, offset, part.Length);
            offset += part.Length;
        }

        private static int CalculatePacketLength(Packet packet)
        {
            int count = 0;

            count += HeaderSize;
            count += NumWordsSize;

            count += WordSizeSize * packet.Words.Count;

            foreach (string word in packet.Words)
            {
                count += word.Length + 1; //null byte
            }

            return count;
        }

        private static byte[] IntToBytes(int value)
        {
            //will use bitconverter because most users will be little endian as well
            byte[] result = BitConverter.GetBytes(value);
            if (!BitConverter.IsLittleEndian)
                Array.Reverse(result);

            return result;
        }
        #endregion

        #region GetPacket Helpers
        private void EnsureBufferHasBytes(int byteCount)
        {
            IEnumerable<byte> newData;

            while (Data.Count < byteCount)
            {
                newData = RequestData();

                foreach (byte b in newData)
                {
                    Data.Enqueue(b);
                }
            }
        }

        private async Task EnsureBufferHasBytesAsync(int byteCount)
        {
            IEnumerable<byte> newData;

            while (Data.Count < byteCount)
            {
                newData = await RequestDataAsync().ConfigureAwait(false);

                foreach (byte b in newData)
                {
                    Data.Enqueue(b);
                }
            }
        }
        private static int BytesToInt(byte[] data, int offset = 0)
        {
            int result;

            //will use bitconverter because most users will be little endian as well
            if (!BitConverter.IsLittleEndian)
                Array.Reverse(data);

            result = BitConverter.ToInt32(data, offset);

            return result;
        }

        private PacketHeader DequeueHeader()
        {
            int header = DequeueInt();
            int length = DequeueInt();
            bool isFromClient = (header & IsFromClientMask) != 0;
            bool isResponse = (header & IsResponseMask) != 0;
            int sequence = header & SequenceMask;

            return new PacketHeader(isFromClient, isResponse, sequence, length);
        }

        private int DequeueInt()
        {
            int result;

            byte[] bytes = DequeueBlock(4);

            result = BytesToInt(bytes);

            return result;
        }

        private byte[] DequeueBlock(int length)
        {
            byte[] bytes = new byte[length];

            for (int i = 0; i < length; i++)
            {
                bytes[i] = Data.Dequeue();
            }

            return bytes;
        }

        private class PacketHeader
        {
            public bool IsFromClient { get; set; }
            public bool IsResponse { get; set; }
            public int Sequence { get; set; }
            public int PacketSize { get; set; }

            public PacketHeader(bool isFromClient, bool isResponse, int sequence, int packetSize)
            {
                IsFromClient = isFromClient;
                IsResponse = isResponse;
                Sequence = sequence;
                PacketSize = packetSize;
            }
        }

        private class WordCollection : IEnumerable<string>
        {
            int WordCount;
            byte[] WordData;

            //sacrificing some overhead for ease of coding
            private IEnumerable<string> GetWords()
            {
                string currentWord;
                int wordLength;
                int offset = 0;

                for (int i = 0; i < WordCount; i++)
                {
                    wordLength = BytesToInt(WordData, offset);
                    offset += WordSizeSize;

                    currentWord = Encoding.UTF8.GetString(WordData, offset, wordLength);

                    offset += wordLength + 1; //null byte

                    yield return currentWord;
                }
            }

            public WordCollection(int wordCount, byte[] wordData)
            {
                WordCount = wordCount;
                WordData = wordData;
            }

            public IEnumerator<string> GetEnumerator()
            {
                return GetWords().GetEnumerator();
            }

            System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
            {
                return GetWords().GetEnumerator();
            }
        }
        #endregion
    }
}
